// Copyright 2022 jeff.li. and/or its affiliates.

#ifndef TBIR_RUNTIME_PIPELINE_OP_KERNEL_H
#define TBIR_RUNTIME_PIPELINE_OP_KERNEL_H

#include <functional>
#include <tuple>
#include <vector>

#include <tbir/pipeline/attributes.h>
#include <tbir/pipeline/global_unique_index.h>
#include <tbir/runtime/native_object_registry.h>

namespace tbir::runtime {

    class TXSession;

    class OpKernel;

    using OpKernelPtr = std::shared_ptr<OpKernel>;

    class OpKernel {
    public:
        OpKernel() = default;

        virtual ~OpKernel() = default;

        OpKernel(const OpKernel &) = default;

        OpKernel(OpKernel &&) = default;

        OpKernel &operator=(const OpKernel &) = default;

        OpKernel &operator=(OpKernel &&) = default;

        template<class T>
        static std::unique_ptr<T> Create(const Dict &config) {
            Attributes attrs = Attributes::FromDict(config);
            auto op = std::make_unique<T>();
            op->Initialize(std::move(attrs));
            return std::move(attrs);
        }

    public:
        TURBO_FORCE_INLINE void CheckArgs(size_t arguments_size, size_t expect_size) const {
            MXCHECK_EQ(arguments_size, expect_size) << "[" << class_name_ << "] Expect " << expect_size
                                                    << " arguments but get " << arguments_size;
        }

        virtual void Init() {
        }

        virtual RTValue Process(PyArgs inputs) const {
            MXCHECK(false) << "[" << class_name_
                           << "] NotImplementedError: The Process method is not implemented";
            return None;
        }

        virtual int Bundle(string_view folder) {
            return 0;
        }

        inline bool HasAttr(string_view key) {
            return attributes_.HasAttr(key);
        }

        template<class U>
        inline U GetAttr(string_view key, const U &default_val = U{}) const {
            return attributes_.GetAttr(key, default_val);
        }

        template<class U>
        inline void SetAttr(string_view key, U &&val) {
            return attributes_.SetAttr(key, std::forward<U>(val));
        }

        String BundlePath(string_view location, string_view folder) const;

    public:
        void Initialize(Attributes attrs);

        string_view ClassName() const {
            return class_name_;
        }

        const String &GetName() const {
            return name_;
        }

        void SetBelongTo(TXSession *belong_to);

        OpKernelPtr GetOpImpl(string_view cls, string_view name);

    protected:
        int device_ = NONE_DEVICE;
        string_view class_name_;
        String name_;
        String resource_path_;
        Attributes attributes_;
        std::vector<OpKernelPtr> sub_ops_;
        TXSession *belong_to_ = nullptr;

        friend class SymbolicExecutor;

        friend class TXSession;

        friend UserDataRef make_op_kernel(string_view, PyArgs args, TXSession *sess);

        friend struct UserDataMutator;
    };

    OpKernelPtr check_get_op_kernel(const UserDataRef &ud);

    OpKernelPtr try_get_op_kernel(const UserDataRef &ud);

    UserDataRef make_userdata(OpKernelPtr op_ptr);

    UserDataRef make_op_kernel(string_view class_name, PyArgs args, TXSession *sess = nullptr);

    inline RTValue op_kernel_call(void *self, PyArgs args) {
        OpKernel *op = reinterpret_cast<OpKernel *>(self);
        return op->Process(args);
    }

}  // namespace tbir::runtime

#define TBIR_REGISTER_NATIVE_OP(ClassName)                                              \
  TBIR_REGISTER_NATIVE_OBJECT(ClassName)                                                \
      .SetConstructor([](::tbir::runtime::PyArgs args) -> std::shared_ptr<void> { \
        MXCHECK(args.size() == 1 && args[0].IsObjectRef<Dict>())                        \
            << "[NativeOp:" << #ClassName                                               \
            << "] only need one dict type arg, but get arg num: " << args.size()        \
            << ", args[0] type: " << args[0].type_name();                               \
        auto op = std::static_pointer_cast<OpKernel>(std::make_shared<ClassName>());    \
        return std::move(op);                                                           \
      })                                                                                \
      .RegisterFunction("__call__", op_kernel_call)

#endif  // TBIR_RUNTIME_PIPELINE_OP_KERNEL_H
